"""
Interface for data managers used in the DIL_Frame.

This module defines the contract that data managers must fulfill to be
compatible with the DIL_Frame learners and components.
"""

from abc import ABC, abstractmethod
from typing import Any, Tuple, Union, Optional, List
from torch.utils.data import Dataset, DataLoader


class DataManagerInterface(ABC):
    """Interface for data management in incremental learning.

    Data managers are responsible for providing access to datasets
    for different tasks, domains, or classes in incremental learning.
    """

    @property
    @abstractmethod
    def nb_tasks(self) -> int:
        """Get total number of tasks."""
        pass

    @abstractmethod
    def get_task_size(self, task_id: int) -> int:
        """Get number of classes or examples in a specific task.

        Args:
            task_id: ID of the task to query

        Returns:
            Number of classes or examples in the task
        """
        pass

    @abstractmethod
    def get_dataset(
        self,
        indices: Union[List[int], int, slice],
        source: str = "train",
        mode: str = "train",
        ret_data: bool = False,
    ) -> Union[Dataset, Tuple[Any, Any, Dataset]]:
        """Get a dataset containing specific classes or tasks.

        Args:
            indices: Class indices, task ID, or slice of classes
            source: Data source ("train", "test", etc.)
            mode: Access mode ("train", "test", etc.)
            ret_data: Whether to return raw data along with the dataset

        Returns:
            Either a dataset or a tuple of (data, targets, dataset)
        """
        pass

    @abstractmethod
    def get_dataloader(
        self,
        dataset: Dataset,
        batch_size: int,
        shuffle: bool = True,
        num_workers: Optional[int] = None,
    ) -> DataLoader:
        """Create a DataLoader from a dataset.

        Args:
            dataset: The dataset to create a loader from
            batch_size: Batch size for the loader
            shuffle: Whether to shuffle the data
            num_workers: Number of worker processes for loading

        Returns:
            A configured DataLoader instance
        """
        pass
